import tensorflow as tf
from matplotlib import pyplot as plt
import onnx
from onnx_tf.backend import prepare
import cv2
import numpy  as np
from timm.models.resnet import resnet50
import logging
import torch
from torch.nn import functional as f
from timm.models import create_model


def read_tensor_from_image_file(file_name,
                                input_height=299,
                                input_width=299):


    input_name = "file_reader"
    output_name = "normalized"

    file_reader = tf.io.read_file(file_name, input_name)
    if file_name.endswith(".png"):
        image_reader = tf.image.decode_png(
            file_reader, channels=3, name="png_reader")
    elif file_name.endswith(".gif"):
        image_reader = tf.squeeze(
            tf.image.decode_gif(file_reader, name="gif_reader"))
    elif file_name.endswith(".bmp"):
        image_reader = tf.image.decode_bmp(file_reader, name="bmp_reader")
    else:
        image_reader = tf.image.decode_jpeg(
            file_reader, channels=3, name="jpeg_reader")
    # float_caster = tf.cast(image_reader, tf.float32)

    float_caster = tf.image.convert_image_dtype(image_reader, dtype=tf.float32)
    float_caster = tf.image.central_crop(float_caster, central_fraction=0.875)

    image = tf.expand_dims(float_caster, 0)
    image = tf.compat.v1.image.resize_bilinear(image, [input_height, input_width])
    image = tf.squeeze(image, [0])
    image = tf.subtract(image, 0.5)
    image = tf.multiply(image, 2.0,name=output_name)


    sess = tf.compat.v1.Session()
    result = sess.run(image)

    return result


def get_graph_from_pb():
    graph = tf.compat.v1.Graph()
    graph_def = tf.compat.v1.GraphDef()
    with open(pb, "rb") as f:
        graph_def.ParseFromString(f.read())
    with graph.as_default():
        tf.import_graph_def(graph_def)
    return graph


def inference_model_torch(model="resnet50"):
    """
    infer the model of pytorch
    :param model: model name
    :return:
    """

    model = create_model(model,num_classes=3,checkpoint_path=torch_model)

    print('Model %s created, param count: %d' %
                     (model, sum([m.numel() for m in model.parameters()])))


    model = model.cuda()
    model.eval()

    with torch.no_grad():
        input_tensor = torch.tensor(img_d).type(torch.float32)
        input = input_tensor.cuda()
        labels = model(input)
        pro = f.softmax(labels)
        print("score: ",labels.cpu().numpy())
        print("probability: ",pro.cpu().numpy())


img = "1.jpg"
pb = "666.pb"
onnx_file = "e:/onnx/pornresnet50Nowithdo_constant_folding.onnx"

torch_model = "e:/onnx/model_best.pth"

img_d = cv2.imread(img)
img_d = cv2.cvtColor(img_d, cv2.COLOR_BGR2RGB)
img_d = cv2.resize(img_d,(224,224))
img_d= img_d.reshape([1,3,224,224])


# 周一测试一下torch转onnx的参数
torch.onnx.export()




def onnx2tf(onnx_file=onnx_file,output_name="777.pb"):

    onnx_model = onnx.load(onnx_file)
    tf_rep = prepare(onnx_model)
    # todo: onnx 可以推理，但转成pb之后暂时无法确定是否可以进行推理
    # output = tf_rep.run(img_d) #Outputs(output=array([[ 5.514201 , -3.9077456, -4.230067 ]], dtype=float32))
    tf_rep.export_graph(output_name)


# todo：暂时无法将pb转为ckpt,onnx
def pb_ckpt():
    with tf.gfile.GFile(pb, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def)
        input = graph.get_tensor_by_name("import/input:0")
        output = graph.get_tensor_by_name("import/add_16:0")

        # sess = tf.Session(graph=graph)
        # saver = tf.train.Saver()
        # saver.save(sess, "666.ckpt")


    # graph = get_graph_from_pb()
    # input = graph.get_operation_by_name("import/input")
    # output = graph.get_operation_by_name("import/InceptionV4/Logits/Predictions")

    with tf.Session(graph=graph) as sess:
        y = sess.run(output,feed_dict={input:img_d})
        print(y)
    #
    #     #
    #     saver.save(sess,"666.ckpt")



if __name__ == '__main__':
    # inference_model(model="resnet50")


    pass



    # graph = get_graph_from_pb()
    # inputnode = graph.get_operation_by_name('import/input')
    # a =graph.get_operations()[-1]
    # print(inputnode)
    # print(a)
